from rich import print
from rich.traceback import install
from typing import (
    TypedDict,
    Optional
)
from train_test import (
    train as train_model,
    valid as valid_model
)
from utils import conf

install()
class InputText(TypedDict):
    text: str
InputTextType = Optional[InputText]

def predict_model(input_text: InputTextType=None) -> None:
    from train_test import predict

    print('[bold gold1]Predicting model...[/]')
    if input_text is None:
        with open(conf.path.original, 'r', encoding='utf-8') as f:
            for line in f:
                input_text = {'text': line.strip()}
                result = predict(input_text)
                print(f'[bold green]result:[/] {result}')
    else:
        result = predict(input_text)
        print(f'[bold green]result:[/] {result}')

def main() -> None:
    input_text: InputTextType = {'text': ''}
    train_model()
    valid_model()
    predict_model()

if __name__ == '__main__':
    main()
